package gov.ornl.stucco.unstructured;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.List;
import java.util.Map;

import edu.stanford.nlp.pipeline.Annotation;
import gov.ornl.stucco.ConfigLoader;
import gov.ornl.stucco.RabbitMQConsumer;
import gov.ornl.stucco.RelationExtractor;
import gov.ornl.stucco.entity.EntityLabeler;
import gov.ornl.stucco.structured.StructuredTransformer;
import gov.pnnl.stucco.doc_service_client.DocServiceClient;
import gov.pnnl.stucco.doc_service_client.DocServiceException;
import gov.ornl.stucco.preprocessors.PreprocessSTIX;
import gov.ornl.stucco.preprocessors.PreprocessSTIX.Vertex;
import gov.ornl.stucco.GraphConstructor;
import gov.ornl.stucco.AlignFactory;
import gov.ornl.stucco.Align;
import gov.ornl.stucco.stix_extractors.StuccoExtractor;

import org.mitre.stix.stix_1.STIXPackage;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.jdom2.Element;

import com.rabbitmq.client.GetResponse;

public class UnstructuredTransformer {
    private static final String PROCESS_NAME = "UNSTRUCTURED";
    private RabbitMQConsumer consumer;
    private DocServiceClient docClient;
    private EntityLabeler entityLabeler;
    private RelationExtractor relationExtractor;
    private PreprocessSTIX preprocessSTIX;
    private GraphConstructor constructGraph;
    private Align alignment;

    private boolean persistent;
    private int sleepTime;

    public UnstructuredTransformer() {
        System.out.println("loading config file from default location");
        ConfigLoader configLoader = new ConfigLoader();
        init(configLoader);
    }

    private void init(ConfigLoader configLoader) {
        Map<String, Object> configMap;
        String exchange = null;
        String queue = null;
        String host = null;
        int port = -1;
        String user = null;
        String password = null;
        String[] bindingKeys = null;
        try {
            configMap = configLoader.getConfig("unstructured_data");
            exchange = String.valueOf(configMap.get("exchange"));
            queue = String.valueOf(configMap.get("queue"));
            host = String.valueOf(configMap.get("host"));
            port = Integer.parseInt(String.valueOf(configMap.get("port")));
            user = String.valueOf(configMap.get("username"));
            password = String.valueOf(configMap.get("password"));
            persistent = Boolean.parseBoolean(String.valueOf(configMap.get("persistent")));
            sleepTime = Integer.parseInt(String.valueOf(configMap.get("emptyQueueSleepTime")));
            List<String> bindings = (List<String>) configMap.get("bindings");
            bindingKeys = new String[bindings.size()];
            bindingKeys = bindings.toArray(bindingKeys);
        } catch (Exception e) {
            System.err.println("Error loading configuration." + e);
            System.exit(-1);
        }
        System.out.println("Config file loaded and parsed");
        try {
            System.out.println("Connecting to rabbitMQ with this info: \nhost: " + host + "\nport: " + port +
                    "\nexchange: " + exchange + "\nqueue: " + queue +
                    "\nuser: " + user + "\npass: " + password);
            consumer = new RabbitMQConsumer(exchange, queue, host, port, user, password, bindingKeys);
            consumer.openQueue();
            entityLabeler = new EntityLabeler();
            relationExtractor = new RelationExtractor();
            preprocessSTIX = new PreprocessSTIX();
            constructGraph = new GraphConstructor();
            alignment = AlignFactory.getAlign();
            configMap = configLoader.getConfig("document_service");
            host = String.valueOf(configMap.get("host"));
            port = Integer.parseInt(String.valueOf(configMap.get("port")));
            docClient = new DocServiceClient(host, port);
        } catch (IOException e) {
            System.err.println("Error initializing Alignment and/or DB connection." + e);
            System.exit(-4);
        }
        System.out.println("Alignment obj, DB connection, and Document service client created.  Initialization complete!");
    }


    public void run() {
        GetResponse response = null;
        boolean fatalError = false;
        do {
            System.out.println("等到消费者信息！");
            try {
                response = consumer.getMessage();
                System.out.println(response);
            } catch (IOException e) {
                System.err.println("mmp消费者死了");
                fatalError = true;
            }
            while (response != null && !fatalError) {
                System.out.println("得到信息了");
                String routingKey = response.getEnvelope().getRoutingKey();
                System.out.println("routingKey:" + routingKey);
                long deliveryTag = response.getEnvelope().getDeliveryTag();
                System.out.println("deliveryTag:" + deliveryTag);
                if (response.getBody() != null) {
                    String message = new String(response.getBody());
                    System.out.println("message:" + message);
                    long timestamp = response.getProps().getTimestamp().getTime();
                    System.out.println("timestamp" + timestamp);
                    boolean contentIncluded = false;
                    Map<String, Object> headerMap = response.getProps().getHeaders();
                    if ((headerMap != null) && (headerMap.containsKey("HasContent"))) {
                        contentIncluded = Boolean.valueOf(String.valueOf(headerMap.get("HasContent")));
                    }
                    System.out.println("Recieved: " + routingKey + " deliveryTag=[" + deliveryTag + "] message- " + message);
                    //Get the document and title from the document server, if necessary
                    String content = message;
                    String title = "";
                    if (!contentIncluded) {
                        String docId = content.trim();
                        System.out.println("Retrieving document content from Document-Service for id '" + docId + "'.");
                        try {
                            JSONObject jsonObject = docClient.fetchExtractedText(docId);
                            content = jsonObject.getString("document");
                            title = jsonObject.getString("title");
                        } catch (DocServiceException e) {
                            System.err.println("Could not fetch document '" + docId + "' from Document-Service." + e);
                        }
                    }
                    //Label the entities/concepts in the document
                    Annotation annotatedDoc = entityLabeler.getAnnotatedDoc(title, content);
                    //Extract the data source name from the routing key
                    String dataSource = routingKey;
                    int index = routingKey.indexOf(PROCESS_NAME.toLowerCase());
                    if (index > -1) {
                        dataSource = routingKey.substring(index + PROCESS_NAME.length());
                        if (dataSource.startsWith(".")) {
                            dataSource = dataSource.substring(1);
                        }
                    }
                    //Construct the subgraph from the concepts and relationships
                    String graphString = relationExtractor.createSubgraph(annotatedDoc, dataSource);
                    if (graphString != null) {
                        try {
                            JSONObject graph = new JSONObject(graphString);
                            StuccoExtractor stuccoExt = new StuccoExtractor(graph);
                            STIXPackage stixPackage = stuccoExt.getStixPackage();
                            Map<String, Vertex> stixElements = preprocessSTIX.normalizeSTIX(stixPackage.toXMLString());
                            graph = constructGraph.constructGraph(stixElements);
                            alignment.load(graph);
                        } catch (RuntimeException e) {
                            System.err.println("Error occurred with routingKey = " + routingKey);
                            System.err.println("										content = " + message);
                            System.err.println("										source = " + dataSource);
                            e.printStackTrace();
                        }
                    }
                    //TODO: Add timestamp into subgraph
                    //Merge subgraph into full knowledge graph
                    //	alignment.load(graph);

                    //Ack the message was processed and can be discarded from the queue
                    try {
                        consumer.messageProcessed(deliveryTag);
                    } catch (IOException e) {
                        System.err.println("Encountered RabbitMQ IO error:" + e);
                        fatalError = true;
                    }
                } else {
                    try {
                        consumer.retryMessage(deliveryTag);
                    } catch (IOException e) {
                        System.err.println("Encountered RabbitMQ IO error:" + e);
                        fatalError = true;
                    }
                }
                //Get next message from queue
                try {
                    response = consumer.getMessage();
                } catch (IOException e) {
                    System.err.println("Encountered RabbitMQ IO error:" + e);
                    fatalError = true;
                }
            }
            //Either the queue is empty, or an error occurred.
            //Either way, sleep for a bit to prevent rapid loop of re-starting.
            try {
                Thread.sleep(sleepTime);
            } catch (InterruptedException consumed) {
                //don't care in this case, exiting anyway.
            }
        } while (persistent && !fatalError);
        try {
            consumer.close();
        } catch (IOException e) {
            System.err.println("Encountered RabbitMQ IO error when closing connection:" + e);
        }
    }


    /**
     * @param args
     */
    public static void main(String[] args) {
        System.setProperty("STUCCO_DB_TYPE", "POSTGRESQL");
        System.setProperty("STUCCO_DB_CONFIG", "postgresql.yml");
        System.setProperty("STUCCO_DB_INDEX_CONFIG", "");
        UnstructuredTransformer unstructProcess = new UnstructuredTransformer();
        unstructProcess.run();
    }
}
